"""
@author:    Guan'an Wang
@contact:   guan.wang0706@gmail.com
"""

import torch
import torch.nn as nn
import copy

import feature_extract.myutils
from feature_extract.registry import MODEL_REGISTRY
from feature_extract.model.light_reid import resnet50ibna
from feature_extract.model.light_reid.bn_head import BNHead

__all__ = ['BaseReIDModel']


@MODEL_REGISTRY.register('BaseReIDModel')
class BaseReIDModel(nn.Module):
    """
    Architecture for ReID Model
    combine backbone, pooling and head modules
    """

    def load_weights(self, path, device):
        check_points = torch.load(path, map_location=device, encoding='iso-8859-1')
        state_dict = utils.remove_prefix(check_points, 'module.')
        self.load_state_dict(state_dict)

    def __init__(self, pretrained=False,class_num=1):
        super(BaseReIDModel, self).__init__()
        self.backbone = resnet50ibna(pretrained=pretrained, last_stride_one=True)
        self.pooling = nn.AdaptiveAvgPool2d(1)
        self.head = BNHead(in_dim=self.backbone.dim, class_num=class_num,
                           classifier={'name': 'circle', 'scale': 64, 'margin': 0.35})
        self.disable_hash()

    def forward(self, x, y=None, fixcnn=False):
        '''
        Args:
            x(torch.tensor): images
            y(torch.tensor): labels, required by circle_softmax, arc_softmax
        '''
        # cnn backbone
        feats_map = self.backbone(x)
        if fixcnn:
            feats_map = feats_map.detach()
        # pooling
        feats_vec = self.pooling(feats_map).squeeze(3).squeeze(2)  # support bs=1
        # head
        res = self.head(feats_vec, y)

        # return
        res['feats_map'] = feats_map
        res['feats_vec'] = feats_vec

        if self.training:
            return res
        else:
            return res[self.test_feats]

    def enable_hash(self):
        self.head.enable_hash()
        self.test_feats = 'binary_feats'

    def disable_hash(self):
        self.head.disable_hash()
        self.test_feats = 'bn_feats'
